
from meantime.models import model_factory
from meantime.dataloaders import dataloader_factory
from meantime.utils import *
from meantime.config import *

from dotmap import DotMap
import sys
import yaml


def count_parameters(model):
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    non_trainable_params = total_params - trainable_params
    return {'Total': total_params, 'Trainable': trainable_params, 'Non-trainable': non_trainable_params}


def init_model(model_name, dataset):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    project_path = os.path.dirname(__file__)
    sys.path.append(project_path)

    template_name = f"train_{model_name}"
    args = yaml.safe_load(open(os.path.join(project_path + '/templates', f'{template_name}.yaml')))
    args = DotMap(args, _dynamic=False)
    args.mode = "val"
    args.add_genre = True
    args.dataset_code = dataset

    dataset_map = {
        "ml-1m": "1m",
        "ml-20m": "20m",
        "beauty": "beauty",
        "game": "game"
    }

    train_loader, val_loader, test_loader = dataloader_factory(args)

    model = model_factory(args)
    args.pretrained_weights = rf"experiments\test\{model_name}_{dataset_map[dataset]}\models\best_model.pth"
    if args.pretrained_weights is not None:
        model.load(args.pretrained_weights)
        print("已成功加载训练好的模型")

    model = model.to(device)
    model.eval()
    return model


if __name__ == '__main__':
    dataset = "ml-20m"
    marank = init_model("marank", dataset)
    # 计算参数
    params = count_parameters(marank)
    print(f"Total parameters: {params['Total']}")
    print(f"Trainable parameters: {params['Trainable']}")
    print(f"Non-trainable parameters: {params['Non-trainable']}")
    print("\n")

    sas = init_model("sas", dataset)
    # 计算参数
    params = count_parameters(sas)
    print(f"Total parameters: {params['Total']}")
    print(f"Trainable parameters: {params['Trainable']}")
    print(f"Non-trainable parameters: {params['Non-trainable']}")
    print("\n")

    tisas = init_model("tisas", dataset)
    # 计算参数
    params = count_parameters(tisas)
    print(f"Total parameters: {params['Total']}")
    print(f"Trainable parameters: {params['Trainable']}")
    print(f"Non-trainable parameters: {params['Non-trainable']}")
    print("\n")

    bert = init_model("bert", dataset)
    # 计算参数
    params = count_parameters(bert)
    print(f"Total parameters: {params['Total']}")
    print(f"Trainable parameters: {params['Trainable']}")
    print(f"Non-trainable parameters: {params['Non-trainable']}")
    print("\n")

    meantime = init_model("meantime", dataset)
    # 计算参数
    params = count_parameters(meantime)
    print(f"Total parameters: {params['Total']}")
    print(f"Trainable parameters: {params['Trainable']}")
    print(f"Non-trainable parameters: {params['Non-trainable']}")
    print("\n")

    modify_meantime = init_model("modify_meantime", dataset)
    # 计算参数
    params = count_parameters(modify_meantime)
    print(f"Total parameters: {params['Total']}")
    print(f"Trainable parameters: {params['Trainable']}")
    print(f"Non-trainable parameters: {params['Non-trainable']}")
